import os
import numpy as np
import pandas as pd
import open3d as o3d
from transforms3d.quaternions import quat2mat


def default_end_link_trans():  # liuyun's
    end_link_trans = np.float32([
        [0, 10, 0],
        [4, -8, 20],
        [-4, -8, 17],
        [-4, -8, 20],
        [4, -8, 17],
        [-2.5, 0, 0],
        [-2.5, 0, 0],
        [-2.5, 0, 0],
        [-2.5, 0, 0],
        [-2.5, 0, 0],
        [2.5, 0, 0],
        [2.5, 0, 0],
        [2.5, 0, 0],
        [2.5, 0, 0],
        [2.5, 0, 0],
    ])
    return end_link_trans


def get_joint_info():
    joint_info = [
        ["Hips", None],
        ["RightUpLeg", "Hips"],
        ["RightLeg", "RightUpLeg"],
        ["RightFoot", "RightLeg"],
        ["LeftUpLeg", "Hips"],
        ["LeftLeg", "LeftUpLeg"],
        ["LeftFoot", "LeftLeg"],
        ["Spine", "Hips"],
        ["Spine1", "Spine"],
        ["Spine2", "Spine1"],
        ["Neck", "Spine2"],
        ["Neck1", "Neck"],
        ["Head", "Neck1"],
        ["RightShoulder", "Spine2"],
        ["RightArm", "RightShoulder"],
        ["RightForeArm", "RightArm"],
        ["RightHand", "RightForeArm"],
        ["RightHandThumb1", "RightHand"],
        ["RightHandThumb2", "RightHandThumb1"],
        ["RightHandThumb3", "RightHandThumb2"],
        ["RightHandIndex", "RightHand"],
        ["RightHandIndex1", "RightHandIndex"],
        ["RightHandIndex2", "RightHandIndex1"],
        ["RightHandIndex3", "RightHandIndex2"],
        ["RightHandMiddle", "RightHand"],
        ["RightHandMiddle1", "RightHandMiddle"],
        ["RightHandMiddle2", "RightHandMiddle1"],
        ["RightHandMiddle3", "RightHandMiddle2"],
        ["RightHandRing", "RightHand"],
        ["RightHandRing1", "RightHandRing"],
        ["RightHandRing2", "RightHandRing1"],
        ["RightHandRing3", "RightHandRing2"],
        ["RightHandPinky", "RightHand"],
        ["RightHandPinky1", "RightHandPinky"],
        ["RightHandPinky2", "RightHandPinky1"],
        ["RightHandPinky3", "RightHandPinky2"],
        ["LeftShoulder", "Spine2"],
        ["LeftArm", "LeftShoulder"],
        ["LeftForeArm", "LeftArm"],
        ["LeftHand", "LeftForeArm"],
        ["LeftHandThumb1", "LeftHand"],
        ["LeftHandThumb2", "LeftHandThumb1"],
        ["LeftHandThumb3", "LeftHandThumb2"],
        ["LeftHandIndex", "LeftHand"],
        ["LeftHandIndex1", "LeftHandIndex"],
        ["LeftHandIndex2", "LeftHandIndex1"],
        ["LeftHandIndex3", "LeftHandIndex2"],
        ["LeftHandMiddle", "LeftHand"],
        ["LeftHandMiddle1", "LeftHandMiddle"],
        ["LeftHandMiddle2", "LeftHandMiddle1"],
        ["LeftHandMiddle3", "LeftHandMiddle2"],
        ["LeftHandRing", "LeftHand"],
        ["LeftHandRing1", "LeftHandRing"],
        ["LeftHandRing2", "LeftHandRing1"],
        ["LeftHandRing3", "LeftHandRing2"],
        ["LeftHandPinky", "LeftHand"],
        ["LeftHandPinky1", "LeftHandPinky"],
        ["LeftHandPinky2", "LeftHandPinky1"],
        ["LeftHandPinky3", "LeftHandPinky2"],
    ]
    return joint_info


def local2world(pos, local2world):
    """
    pos: shape = (3,)
    把局部link坐标系的3D坐标转化到全局坐标系
    """
    pos = np.float32(pos)
    assert pos.shape == (3,)
    pos = np.hstack((pos, [1.0])).reshape(4, 1)
    pos = local2world @ pos
    pos = pos[:3].reshape(3)
    return pos


def add_end_link(joint_3Dpos, joint_acc_pose, end_link_trans):
    assert len(joint_3Dpos) == 59

    joint_3Dpos.append(local2world(end_link_trans[0], joint_acc_pose["Head"]))
    joint_3Dpos.append(local2world(end_link_trans[1], joint_acc_pose["RightFoot"]))
    joint_3Dpos.append(local2world(end_link_trans[2], joint_acc_pose["RightFoot"]))
    joint_3Dpos.append(local2world(end_link_trans[3], joint_acc_pose["LeftFoot"]))
    joint_3Dpos.append(local2world(end_link_trans[4], joint_acc_pose["LeftFoot"]))
    joint_3Dpos.append(local2world(end_link_trans[5], joint_acc_pose["RightHandThumb3"]))
    joint_3Dpos.append(local2world(end_link_trans[6], joint_acc_pose["RightHandIndex3"]))
    joint_3Dpos.append(local2world(end_link_trans[7], joint_acc_pose["RightHandMiddle3"]))
    joint_3Dpos.append(local2world(end_link_trans[8], joint_acc_pose["RightHandRing3"]))
    joint_3Dpos.append(local2world(end_link_trans[9], joint_acc_pose["RightHandPinky3"]))
    joint_3Dpos.append(local2world(end_link_trans[10], joint_acc_pose["LeftHandThumb3"]))
    joint_3Dpos.append(local2world(end_link_trans[11], joint_acc_pose["LeftHandIndex3"]))
    joint_3Dpos.append(local2world(end_link_trans[12], joint_acc_pose["LeftHandMiddle3"]))
    joint_3Dpos.append(local2world(end_link_trans[13], joint_acc_pose["LeftHandRing3"]))
    joint_3Dpos.append(local2world(end_link_trans[14], joint_acc_pose["LeftHandPinky3"]))

    assert len(joint_3Dpos) == 74
    return joint_3Dpos


def bvh2joint(joint_rawdata, end_link_trans=None, return_local_rot=False):
    """
    joint_rawdata: VTS broadcasted data
    end_link_trans: 人体15个叶节点在T-POSE下相对于其父节点的3D偏移, 顺序: [下颚到鼻尖的垂直距离, 右脚(大脚趾,小脚趾)尖到脚腕, 左脚(大脚趾,小脚趾)尖到脚腕, 右手5个指尖长(从thumb到pinky), 左手5个指尖长(从thumb到pinky)], 单位: cm
    头: [0, 下颚到鼻尖的垂直距离, 0]
    左右脚: [(左脚-+,右脚+1)一半脚宽, -脚高, 大/小脚趾尖到脚腕在向前方向上的距离]
    右指尖长: [-指尖长, 0, 0]
    左指尖长: [指尖长, 0, 0]

    return: shape = (N, 59/74, 3)
    """

    if not end_link_trans is None:
        assert end_link_trans.shape == (15, 3)

    joint_info = get_joint_info()
    joint_ids = {}
    for i in range(59):
        joint_ids[joint_info[i][0]] = i

    joint_data = []
    local_rot_data = []

    joint_acc_pose = {}
    for i in range(59):
        pos = joint_rawdata[i]["position"]
        # compensate VTS system error
        if (i == 16) or (i == 39):
            pos[2] += 2  # unit: cm
        if not joint_info[i][1] is None:
            pos = local2world(pos, joint_acc_pose[joint_info[i][1]])
        joint_data.append(pos)
        pose = np.eye(4)
        pose[:3, :3] = quat2mat(joint_rawdata[i]["orientation"])
        pose[:3, 3] = joint_rawdata[i]["position"]
        local_rot_data.append(pose)
        if not joint_info[i][1] is None:
            pose = joint_acc_pose[joint_info[i][1]] @ pose
        joint_acc_pose[joint_info[i][0]] = pose
    
    if not end_link_trans is None:  # 在原先的59个joint的世界系3D坐标后面拼接15个末端的世界系3D坐标，顺序: [头顶, 右脚2个顶点, 左脚2个顶点, 右手5个指尖(从thumb到pinky), 左手5个指尖(从thumb到pinky)]
        joint_data = add_end_link(joint_data, joint_acc_pose, end_link_trans)

    joint_data = np.float32(joint_data) / 100  # unit: m
    local_rot_data = np.float32(local_rot_data)
    local_rot_data[:, :3, 3] /= 100  # unit: m
    if return_local_rot:
        return joint_data, local_rot_data
    else:
        return joint_data
    
    
def localpose_to_headrot(localrot):
    """
    localrot: None or (59, 4, 4)
    return None or (3, 3)
    """
    
    if localrot is None:
        return None
    
    # Hips @ Spine @ Spine1 @ Spine2 @ Neck @ Neck1 @ Head
    head_rot =  localrot[0, :3, :3] @ localrot[7, :3, :3] @ localrot[8, :3, :3] @ localrot[9, :3, :3] @ localrot[10, :3, :3] @ localrot[11, :3, :3] @ localrot[12, :3, :3]
    return head_rot


if __name__ == "__main__":
    joints_data = None
    end_link_trans = default_end_link_trans()
    joint_data, local_rot = bvh2joint(joints_data, end_link_trans=end_link_trans, return_local_rot=True)
